import pygame
import numpy as np
import math
import random
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from scipy.spatial.transform import Rotation as R
from collections import deque

# 初始化pygame
pygame.init()
WIDTH, HEIGHT = 1200, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("无人机峡谷导航与定位融合仿真")
clock = pygame.time.Clock()
font = pygame.font.SysFont('SimHei', 16)
title_font = pygame.font.SysFont('SimHei', 24, bold=True)

# 颜色定义
BACKGROUND = (10, 20, 30)
DRONE_COLOR = (70, 130, 180)
TRUE_PATH_COLOR = (50, 205, 50)
FUSION_PATH_COLOR = (220, 20, 60)
IMU_PATH_COLOR = (255, 215, 0)
GPS_PATH_COLOR = (30, 144, 255)
WALL_COLOR = (100, 100, 120)
TEXT_COLOR = (220, 220, 220)
SENSOR_COLOR = (70, 130, 180, 150)

# 仿真参数
FPS = 60
SIM_SPEED = 1.0  # 仿真速度
DT = 0.1 * SIM_SPEED  # 时间步长

# 峡谷参数
CANYON_WIDTH = 50
CANYON_LENGTH = 300
CANYON_HEIGHT = 100
CANYON_ANGLE = math.radians(45)  # 峡谷走向角度

# 无人机初始状态 - 使用浮点数
INIT_POS = np.array([50.0, 50.0, 80.0])  # 东, 北, 天 (m)
INIT_VEL = np.array([5.0, 5.0, 0.0])  # m/s (对应8m/s速度)
INIT_YAW = math.radians(45)  # 偏航角

def to_mpl_color(rgb_tuple):
    return tuple(c/255 for c in rgb_tuple)

# 传感器配置
class Sensor:
    def __init__(self):
        # IMU参数
        self.imu_freq = 200  # Hz
        self.acc_bias = np.array([0.01, -0.005, 0.003])  # m/s²
        self.gyro_bias = np.array([0.001, -0.002, 0.0005])  # rad/s
        self.acc_noise = 0.02  # m/s²
        self.gyro_noise = 0.005  # rad/s

        # VO参数
        self.vo_freq = 30  # Hz
        self.vo_scale_drift = 0.0001  # 尺度漂移率
        self.vo_noise = 0.05  # m

        # LiDAR参数
        self.lidar_freq = 10  # Hz
        self.lidar_noise = 0.1  # m
        self.lidar_angle_noise = math.radians(0.1)  # rad

        # 气压计参数
        self.baro_freq = 10  # Hz
        self.baro_drift = 0.0001  # m/s
        self.baro_noise = 0.3  # m
        self.baro_offset = 0.0

        # UWB参数
        self.uwb_freq = 1  # Hz
        self.uwb_noise = 0.5  # m
        self.uwb_anchor1 = np.array([0.0, 0.0, 80.0])  # 入口锚点
        self.uwb_anchor2 = np.array([CANYON_LENGTH * math.cos(CANYON_ANGLE),
                                     CANYON_LENGTH * math.sin(CANYON_ANGLE), 80.0])  # 出口锚点

        # 计时器
        self.imu_timer = 0.0
        self.vo_timer = 0.0
        self.lidar_timer = 0.0
        self.baro_timer = 0.0
        self.uwb_timer = 0.0

        # 状态变量
        self.last_vo_position = INIT_POS.copy()
        self.vo_scale = 1.0


# 卡尔曼滤波器实现
class ExtendedKalmanFilter:
    def __init__(self):
        # 状态向量: [x, y, z, vx, vy, vz, roll, pitch, yaw]
        self.state = np.zeros(9, dtype=np.float64)
        self.state[0:3] = INIT_POS  # 位置
        self.state[3:6] = INIT_VEL  # 速度
        self.state[8] = INIT_YAW  # 偏航角

        # 状态协方差矩阵
        self.P = np.diag([1.0, 1.0, 1.0, 0.1, 0.1, 0.1, 0.01, 0.01, 0.01])

        # 过程噪声协方差
        self.Q = np.diag([0.01, 0.01, 0.01, 0.05, 0.05, 0.05, 0.001, 0.001, 0.001])

        # 观测噪声协方差
        self.R_gps = np.diag([2.0, 2.0, 3.0, 0.5, 0.5, 0.5])
        self.R_vo = np.diag([0.5, 0.5, 0.5])
        self.R_baro = np.array([[0.5]])
        self.R_uwb = np.array([[1.0]])
        self.R_lidar = np.array([[0.2]])

        # 重力加速度
        self.g = 9.81

        # 其他协方差矩阵初始化...
        self.last_vo_position = INIT_POS.copy()  # 关键修复

        # 传感器数据缓存
        self.last_imu_data = None
        self.last_vo_delta = np.zeros(3, dtype=np.float64)

    def predict(self, imu_data, dt):
        # 解析IMU数据 [ax, ay, az, gx, gy, gz]
        acc = imu_data[0:3].astype(np.float64)
        gyro = imu_data[3:6].astype(np.float64)

        # 保存IMU数据用于后续更新
        self.last_imu_data = imu_data

        # 获取当前状态
        x, y, z, vx, vy, vz, roll, pitch, yaw = self.state

        # 姿态更新 (使用四元数)
        rotation = R.from_euler('xyz', [roll, pitch, yaw])
        delta_angle = gyro * dt
        rotation = rotation * R.from_rotvec(delta_angle)
        roll, pitch, yaw = rotation.as_euler('xyz')

        # 将加速度从机体坐标系转换到导航坐标系
        acc_nav = rotation.apply(acc)
        acc_nav[2] -= self.g  # 补偿重力

        # 位置和速度更新
        vx += acc_nav[0] * dt
        vy += acc_nav[1] * dt
        vz += acc_nav[2] * dt

        x += vx * dt + 0.5 * acc_nav[0] * dt ** 2
        y += vy * dt + 0.5 * acc_nav[1] * dt ** 2
        z += vz * dt + 0.5 * acc_nav[2] * dt ** 2

        # 更新状态
        self.state = np.array([x, y, z, vx, vy, vz, roll, pitch, yaw], dtype=np.float64)

        # 更新状态转移雅可比矩阵 F
        F = np.eye(9, dtype=np.float64)
        F[0, 3] = dt
        F[1, 4] = dt
        F[2, 5] = dt

        # 更新协方差矩阵
        self.P = F @ self.P @ F.T + self.Q

    def update_gps(self, gps_data):
        # gps_data: [x, y, z, vx, vy, vz]
        H = np.zeros((6, 9), dtype=np.float64)
        H[0:6, 0:6] = np.eye(6)

        # 计算卡尔曼增益
        S = H @ self.P @ H.T + self.R_gps
        K = self.P @ H.T @ np.linalg.inv(S)

        # 更新状态
        y = gps_data - H @ self.state
        self.state += K @ y

        # 更新协方差
        I = np.eye(9, dtype=np.float64)
        self.P = (I - K @ H) @ self.P

    def update_vo(self, vo_delta):
        # VO提供的是相对位移增量
        H = np.zeros((3, 9), dtype=np.float64)
        H[0:3, 0:3] = np.eye(3)

        # 计算卡尔曼增益
        S = H @ self.P @ H.T + self.R_vo
        K = self.P @ H.T @ np.linalg.inv(S)

        # 计算预测位移
        predicted_delta = self.state[0:3] - self.last_vo_position
        self.last_vo_position = self.state[0:3].copy()

        # 更新状态
        y = vo_delta - predicted_delta
        self.state += K @ y

        # 更新协方差
        I = np.eye(9, dtype=np.float64)
        self.P = (I - K @ H) @ self.P

    def update_baro(self, z):
        H = np.zeros((1, 9), dtype=np.float64)
        H[0, 2] = 1

        # 计算卡尔曼增益
        S = H @ self.P @ H.T + self.R_baro
        K = self.P @ H.T @ np.linalg.inv(S)

        # 更新状态
        y = z - self.state[2]
        self.state += K.reshape(-1) * y

        # 更新协方差
        I = np.eye(9, dtype=np.float64)
        self.P = (I - np.outer(K, H)) @ self.P

    def update_uwb(self, distance, anchor_pos):
        # 计算预测距离
        dx = self.state[0] - anchor_pos[0]
        dy = self.state[1] - anchor_pos[1]
        dz = self.state[2] - anchor_pos[2]
        predicted_dist = math.sqrt(dx ** 2 + dy ** 2 + dz ** 2)

        # 计算雅可比矩阵
        H = np.zeros((1, 9), dtype=np.float64)
        H[0, 0] = dx / predicted_dist
        H[0, 1] = dy / predicted_dist
        H[0, 2] = dz / predicted_dist

        # 计算卡尔曼增益
        S = H @ self.P @ H.T + self.R_uwb
        K = self.P @ H.T @ np.linalg.inv(S)

        # 更新状态
        y = distance - predicted_dist
        self.state += K.reshape(-1) * y

        # 更新协方差
        I = np.eye(9, dtype=np.float64)
        self.P = (I - np.outer(K, H)) @ self.P

    def update_lidar(self, measured_distance, expected_distance, wall_normal):
        # 计算距离残差
        distance_residual = measured_distance - expected_distance

        # 计算雅可比矩阵 (距离相对于位置的变化率)
        H = np.zeros((1, 9), dtype=np.float64)
        H[0, 0:3] = wall_normal

        # 计算卡尔曼增益
        S = H @ self.P @ H.T + self.R_lidar
        K = self.P @ H.T @ np.linalg.inv(S)

        # 更新状态
        self.state += K.reshape(-1) * distance_residual

        # 更新协方差
        I = np.eye(9, dtype=np.float64)
        self.P = (I - np.outer(K, H)) @ self.P


# 无人机仿真类
class DroneSimulator:
    def __init__(self):
        self.reset()

    def reset(self):
        # 真实状态 - 确保使用浮点数
        self.true_pos = INIT_POS.copy().astype(np.float64)
        self.true_vel = INIT_VEL.copy().astype(np.float64)
        self.true_yaw = INIT_YAW
        self.true_roll = 0.0
        self.true_pitch = 0.0

        # 传感器
        self.sensor = Sensor()

        # 滤波器
        self.ekf = ExtendedKalmanFilter()

        # 路径记录
        self.true_path = []
        self.fusion_path = []
        self.imu_path = []
        self.gps_path = []

        # 位置估算 (纯IMU)
        self.imu_pos = INIT_POS.copy().astype(np.float64)
        self.imu_vel = INIT_VEL.copy().astype(np.float64)
        self.imu_yaw = INIT_YAW

        # 时间
        self.time = 0.0
        self.gps_lost = False
        self.gps_recovered = False

        # 性能指标
        self.position_errors = []
        self.gps_status = "正常"

        # 添加初始位置
        self.record_paths()

    def record_paths(self):
        self.true_path.append(self.true_pos.copy())
        self.fusion_path.append(self.ekf.state[0:3].copy())
        self.imu_path.append(self.imu_pos.copy())

        # 仅当GPS可用时记录GPS路径
        if not self.gps_lost or self.gps_recovered:
            self.gps_path.append(self.ekf.state[0:3].copy())

    def update(self, dt):
        self.time += dt

        # 更新真实状态 (简单的直线运动模型)
        self.true_pos = self.true_pos + self.true_vel * dt
        self.true_yaw = math.atan2(self.true_vel[1], self.true_vel[0])

        # 添加一些轻微的运动变化
        if 10 < self.time < 40:
            # 在峡谷中轻微偏航
            self.true_vel[0] += random.uniform(-0.01, 0.01)
            self.true_vel[1] += random.uniform(-0.01, 0.01)

        # 检查GPS状态
        if 10 < self.time < 40:
            self.gps_lost = True
            self.gps_status = "信号丢失"
        elif self.time >= 40 and not self.gps_recovered:
            self.gps_lost = False
            self.gps_recovered = True
            self.gps_status = "已恢复"

        # ================= 传感器数据生成 =================
        # IMU数据生成 (高频率)
        self.sensor.imu_timer += dt
        imu_dt = 1.0 / self.sensor.imu_freq
        while self.sensor.imu_timer >= imu_dt:
            # 加速度 (机体坐标系)
            acc_body = np.array([
                0.0,  # 轻微加速度变化
                0.0,
                -9.81  # 重力
            ], dtype=np.float64)

            # 陀螺仪 (机体坐标系)
            gyro_body = np.array([
                random.uniform(-0.01, 0.01),
                random.uniform(-0.01, 0.01),
                0.0
            ], dtype=np.float64)

            # 添加噪声和偏差
            acc_noise = np.random.normal(0, self.sensor.acc_noise, 3).astype(np.float64)
            gyro_noise = np.random.normal(0, self.sensor.gyro_noise, 3).astype(np.float64)

            imu_data = np.concatenate((
                acc_body + self.sensor.acc_bias + acc_noise,
                gyro_body + self.sensor.gyro_bias + gyro_noise
            ))

            # 纯IMU推算 (用于对比)
            rotation_imu = R.from_euler('xyz', [0, 0, self.imu_yaw])
            acc_nav_imu = rotation_imu.apply(acc_body + self.sensor.acc_bias)
            acc_nav_imu[2] += 9.81  # 补偿重力

            self.imu_vel = self.imu_vel + acc_nav_imu * imu_dt
            self.imu_pos = self.imu_pos + self.imu_vel * imu_dt
            self.imu_yaw += gyro_body[2] * imu_dt

            # 更新滤波器
            self.ekf.predict(imu_data, imu_dt)

            self.sensor.imu_timer -= imu_dt

        # VO数据生成
        self.sensor.vo_timer += dt
        if self.sensor.vo_timer >= 1.0 / self.sensor.vo_freq:
            # VO提供的相对位移
            true_delta = self.true_vel * (1.0 / self.sensor.vo_freq)

            # 添加噪声和尺度漂移
            noise = np.random.normal(0, self.sensor.vo_noise, 3).astype(np.float64)
            self.sensor.vo_scale += self.sensor.vo_scale_drift * dt
            vo_delta = true_delta * self.sensor.vo_scale + noise

            # 更新滤波器
            if self.gps_lost:
                self.ekf.update_vo(vo_delta)

            self.sensor.vo_timer = 0

        # LiDAR数据生成
        self.sensor.lidar_timer += dt
        if self.sensor.lidar_timer >= 1.0 / self.sensor.lidar_freq:
            # 计算到峡谷两侧墙壁的距离
            canyon_direction = np.array([math.cos(CANYON_ANGLE), math.sin(CANYON_ANGLE), 0], dtype=np.float64)
            canyon_normal = np.array([-math.sin(CANYON_ANGLE), math.cos(CANYON_ANGLE), 0], dtype=np.float64)

            # 无人机到中心线的投影距离
            # 修改第402行
            center_proj = np.dot(np.append(self.true_pos[:2], 0), canyon_normal)  # 添加z=0分量
            distance_to_left = CANYON_WIDTH / 2 - center_proj
            distance_to_right = CANYON_WIDTH / 2 + center_proj

            # 添加噪声
            noise_left = random.gauss(0, self.sensor.lidar_noise)
            noise_right = random.gauss(0, self.sensor.lidar_noise)

            measured_left = max(0.1, distance_to_left + noise_left)
            measured_right = max(0.1, distance_to_right + noise_right)

            # 更新滤波器 (使用左侧墙壁作为示例)
            if self.gps_lost:
                # 理想情况下，无人机应该距离墙壁25米
                self.ekf.update_lidar(measured_left, CANYON_WIDTH / 2, canyon_normal)

            self.sensor.lidar_timer = 0

        # 气压计数据生成
        self.sensor.baro_timer += dt
        if self.sensor.baro_timer >= 1.0 / self.sensor.baro_freq:
            # 添加漂移和噪声
            self.sensor.baro_offset += random.uniform(-0.01, 0.01)
            baro_z = self.true_pos[2] + self.sensor.baro_offset + random.gauss(0, self.sensor.baro_noise)

            # 更新滤波器
            self.ekf.update_baro(baro_z)

            self.sensor.baro_timer = 0

        # UWB数据生成
        self.sensor.uwb_timer += dt
        if self.sensor.uwb_timer >= 1.0 / self.sensor.uwb_freq:
            # 计算到锚点的真实距离
            dist1 = np.linalg.norm(self.true_pos - self.sensor.uwb_anchor1)
            dist2 = np.linalg.norm(self.true_pos - self.sensor.uwb_anchor2)

            # 添加噪声
            noise1 = random.gauss(0, self.sensor.uwb_noise)
            noise2 = random.gauss(0, self.sensor.uwb_noise)

            measured_dist1 = dist1 + noise1
            measured_dist2 = dist2 + noise2

            # 更新滤波器 (使用入口锚点作为示例)
            if self.gps_lost:
                self.ekf.update_uwb(measured_dist1, self.sensor.uwb_anchor1)

            self.sensor.uwb_timer = 0

        # 如果有GPS信号，进行GPS更新
        if not self.gps_lost:
            gps_noise = np.array([
                random.gauss(0, 1.5),
                random.gauss(0, 1.5),
                random.gauss(0, 3.0),
                random.gauss(0, 0.5),
                random.gauss(0, 0.5),
                random.gauss(0, 0.5)
            ], dtype=np.float64)
            gps_data = np.concatenate((self.true_pos, self.true_vel)) + gps_noise
            self.ekf.update_gps(gps_data)

        # 记录路径
        self.record_paths()

        # 计算位置误差
        fusion_pos = self.ekf.state[0:3]
        error = np.linalg.norm(fusion_pos - self.true_pos)
        self.position_errors.append(error)

        return True

    def draw(self, screen):
        # 绘制背景
        screen.fill(BACKGROUND)

        # 绘制标题
        title = title_font.render("无人机峡谷导航与定位融合仿真", True, (255, 215, 0))
        screen.blit(title, (WIDTH // 2 - title.get_width() // 2, 10))

        # 绘制峡谷
        canyon_start = (50, HEIGHT - 150)
        canyon_end = (canyon_start[0] + CANYON_LENGTH * 0.8, canyon_start[1] - CANYON_LENGTH * 0.5)

        # 峡谷墙壁
        pygame.draw.line(screen, WALL_COLOR, canyon_start, canyon_end, CANYON_WIDTH)
        pygame.draw.line(screen, WALL_COLOR,
                         (canyon_start[0], canyon_start[1] - CANYON_WIDTH),
                         (canyon_end[0], canyon_end[1] - CANYON_WIDTH), 3)

        # 绘制路径 (缩放和平移)
        scale = 0.8
        offset_x, offset_y = 50, HEIGHT - 150

        # 真实路径
        if len(self.true_path) > 1:
            true_points = []
            for pos in self.true_path:
                x = offset_x + pos[0] * scale
                y = offset_y - pos[1] * scale
                true_points.append((x, y))
            pygame.draw.lines(screen, TRUE_PATH_COLOR, False, true_points, 2)

        # 融合估计路径
        if len(self.fusion_path) > 1:
            fusion_points = []
            for pos in self.fusion_path:
                x = offset_x + pos[0] * scale
                y = offset_y - pos[1] * scale
                fusion_points.append((x, y))
            pygame.draw.lines(screen, FUSION_PATH_COLOR, False, fusion_points, 2)

        # IMU路径
        if len(self.imu_path) > 1:
            imu_points = []
            for pos in self.imu_path:
                x = offset_x + pos[0] * scale
                y = offset_y - pos[1] * scale
                imu_points.append((x, y))
            pygame.draw.lines(screen, IMU_PATH_COLOR, False, imu_points, 1)

        # GPS路径
        if len(self.gps_path) > 1:
            gps_points = []
            for pos in self.gps_path:
                x = offset_x + pos[0] * scale
                y = offset_y - pos[1] * scale
                gps_points.append((x, y))
            pygame.draw.lines(screen, GPS_PATH_COLOR, False, gps_points, 1)

        # 绘制当前无人机位置
        true_x = offset_x + self.true_pos[0] * scale
        true_y = offset_y - self.true_pos[1] * scale
        pygame.draw.circle(screen, TRUE_PATH_COLOR, (int(true_x), int(true_y)), 6)

        fusion_x = offset_x + self.ekf.state[0] * scale
        fusion_y = offset_y - self.ekf.state[1] * scale
        pygame.draw.circle(screen, FUSION_PATH_COLOR, (int(fusion_x), int(fusion_y)), 6)

        # 绘制无人机图示
        drone_size = 15
        drone_pos = (int(true_x), int(true_y))
        pygame.draw.circle(screen, DRONE_COLOR, drone_pos, drone_size, 2)

        # 绘制方向指示
        direction_x = drone_pos[0] + math.cos(self.true_yaw) * drone_size
        direction_y = drone_pos[1] - math.sin(self.true_yaw) * drone_size
        pygame.draw.line(screen, DRONE_COLOR, drone_pos, (direction_x, direction_y), 2)

        # 绘制传感器状态
        sensor_rect = pygame.Rect(WIDTH - 300, 50, 280, 250)
        pygame.draw.rect(screen, (40, 40, 60), sensor_rect, border_radius=10)
        pygame.draw.rect(screen, (80, 80, 120), sensor_rect, 2, border_radius=10)

        sensor_title = font.render("传感器状态", True, TEXT_COLOR)
        screen.blit(sensor_title, (sensor_rect.x + 10, sensor_rect.y + 10))

        # 传感器状态文本
        y_offset = 40
        sensors = [
            f"GPS状态: {self.gps_status}",
            f"IMU: 正常 (200Hz)",
            f"视觉里程计: 正常 (30Hz)",
            f"激光雷达: 正常 (10Hz)",
            f"气压计: 正常 (10Hz)",
            f"UWB: 正常 (1Hz)",
            f"融合定位更新: {self.ekf.state[0]:.1f}, {self.ekf.state[1]:.1f}, {self.ekf.state[2]:.1f}",
            f"真实位置: {self.true_pos[0]:.1f}, {self.true_pos[1]:.1f}, {self.true_pos[2]:.1f}",
            f"位置误差: {self.position_errors[-1] if self.position_errors else 0:.2f} m",
            f"仿真时间: {self.time:.1f} s"
        ]

        for sensor in sensors:
            text = font.render(sensor, True, TEXT_COLOR)
            screen.blit(text, (sensor_rect.x + 15, sensor_rect.y + y_offset))
            y_offset += 25

        # 绘制性能图表
        if len(self.position_errors) > 0:
            self.draw_performance_graph(screen)

        # 绘制图例
        self.draw_legend(screen)

    def draw_performance_graph(self, screen):
        # 全局字体设置
        plt.rcParams['font.family'] = 'Microsoft YaHei'  # 更通用的微软雅黑
        plt.rcParams['axes.unicode_minus'] = False
        # 创建位置误差图表
        fig, ax = plt.subplots(figsize=(4, 2.5), dpi=80)
        ax.plot(self.position_errors, color=to_mpl_color(FUSION_PATH_COLOR))
        ax.set_title('位置估计误差', fontsize=10)
        ax.set_xlabel('时间步', fontsize=8)
        ax.set_ylabel('误差 (m)', fontsize=8)
        ax.grid(True, linestyle='--', alpha=0.6)
        ax.tick_params(axis='both', which='major', labelsize=7)
        fig.tight_layout()

        # 将图表渲染到pygame表面
        canvas = FigureCanvasAgg(fig)
        canvas.draw()
        renderer = canvas.get_renderer()
        raw_data = renderer.tostring_rgb()
        size = canvas.get_width_height()

        # 创建pygame表面
        surf = pygame.image.fromstring(raw_data, size, "RGB")
        screen.blit(surf, (WIDTH - 300, HEIGHT - 200))

        plt.close(fig)

    def draw_legend(self, screen):
        legend_rect = pygame.Rect(50, 50, 200, 150)
        pygame.draw.rect(screen, (40, 40, 60), legend_rect, border_radius=10)
        pygame.draw.rect(screen, (80, 80, 120), legend_rect, 2, border_radius=10)

        legend_title = font.render("图例说明", True, TEXT_COLOR)
        screen.blit(legend_title, (legend_rect.x + 10, legend_rect.y + 10))

        # 图例项
        y_offset = 40
        items = [
            (TRUE_PATH_COLOR, "真实路径"),
            (FUSION_PATH_COLOR, "融合定位路径"),
            (IMU_PATH_COLOR, "纯IMU推算路径"),
            (GPS_PATH_COLOR, "GPS定位路径"),
            (WALL_COLOR, "峡谷墙壁")
        ]

        for color, text in items:
            pygame.draw.circle(screen, color, (legend_rect.x + 20, legend_rect.y + y_offset), 6)
            text_surf = font.render(text, True, TEXT_COLOR)
            screen.blit(text_surf, (legend_rect.x + 35, legend_rect.y + y_offset - 8))
            y_offset += 25


# 主游戏循环
def main():
    simulator = DroneSimulator()
    running = True
    paused = False

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    paused = not paused
                elif event.key == pygame.K_r:
                    simulator = DroneSimulator()  # 重置仿真

        if not paused:
            simulator.update(DT)

        simulator.draw(screen)

        # 显示控制提示
        controls = font.render("空格键: 暂停/继续  |  R键: 重置仿真", True, TEXT_COLOR)
        screen.blit(controls, (WIDTH // 2 - controls.get_width() // 2, HEIGHT - 30))

        pygame.display.flip()
        clock.tick(FPS)

    pygame.quit()


if __name__ == "__main__":
    main()